package main

import (
	"fmt"
	"log"
	"os"
	"os/signal"

	"github.com/Shopify/sarama"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"github.com/linkedin/goavro"
)

func main() {
	codec, err := goavro.NewCodec(`
        {"namespace": "com.lee",
			"type": "record",
			"name": "TCP",
			"fields": [
				{"name": "Length",  "type":   "int"  },
				{"name": "Time",  "type":   "long"  },
				{"name": "SrcMAC",  "type":   "string"  },
				{"name": "DstMAC",  "type":   "string"  },
				{"name": "NetworkProtocol",  "type":   "string"  },
				{"name": "IHL",  "type":   "int"  },
				{"name": "TOS",  "type":   "int"  },
				{"name": "Id",  "type":   "int"  },
				{"name": "Flags",  "type":   "string"  },
				{"name": "FragOffset",  "type":   "int"  },
				{"name": "TTL",  "type":   "int"  },
				{"name": "SrcIP",  "type":   "string"  },
				{"name": "DstIP",  "type":   "string"  },
				{"name": "TransportProtol",  "type":   "string"  },
				{"name": "SrcPort",  "type":   "int"  },
				{"name": "DstPort",  "type":   "int"  },
				{"name": "Seq",  "type":   "long"  },
				{"name": "Ack",  "type":   "long"  },
				{"name": "FIN",  "type":   "boolean"  },
				{"name": "SYN",  "type":   "boolean"  },
				{"name": "RST",  "type":   "boolean"  },
				{"name": "PSH",  "type":   "boolean"  },
				{"name": "ACK",  "type":   "boolean"  },
				{"name": "URG",  "type":   "boolean"  },
				{"name": "ECE",  "type":   "boolean"  },
				{"name": "CWR",  "type":   "boolean"  },
				{"name": "NS",  "type":   "boolean"  },
				{"name": "Window",  "type":   "int"  },
				{"name": "CheckSum",  "type":   "int"  }
			]
		   }`)
	if err != nil {
		panic(err)
	}
	path := "/Users/liruopeng/Downloads/test2.pcap"

	// Trap SIGINT to trigger a shutdown.
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt)

	var (
		successes, errors int
	)

	producer := initProducer("10.36.8.128:9092")

	go func() {
		for range (*producer).Successes() {
			fmt.Println(successes)
			successes++
		}
	}()

	go func() {
		for err := range (*producer).Errors() {
			log.Println(err)
			errors++
		}
	}()

	sendToKafka(path, producer, codec)

WaitLoop:
	for {
		select {
		case <-signals:
			(*producer).AsyncClose()
			break WaitLoop
		}
	}
	fmt.Println(successes, errors)
}

func sendToKafka(path string, producer *sarama.AsyncProducer, codec *goavro.Codec) {
	handle, err := pcap.OpenOffline(path)
	if err != nil {
		log.Fatal(err)
	}
	defer handle.Close()
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())

	for packet := range packetSource.Packets() {
		packetMap := GetAvroMap(packet)
		if len(packetMap) < 29 {
			continue
		}
		binary, err := codec.BinaryFromNative(nil, packetMap)
		if err != nil {
			panic(err)
		}
		fmt.Println(binary)
		message := &sarama.ProducerMessage{Topic: "avroTest2", Value: sarama.ByteEncoder(binary)}

		(*producer).Input() <- message

	}
}

func initProducer(url string) *sarama.AsyncProducer {
	config := sarama.NewConfig()
	config.Producer.MaxMessageBytes = 1000000
	// config.Producer.Return.Successes = true
	producer, err := sarama.NewAsyncProducer([]string{url}, config)
	if err != nil {
		panic(err)
	}
	return &producer
}

func GetAvroMap(packet gopacket.Packet) map[string]interface{} {

	data := map[string]interface{}{}
	data["Length"] = len(packet.Data())
	data["Time"] = packet.Metadata().Timestamp.Unix()

	// 网络接口层
	ethernetLayer := packet.Layer(layers.LayerTypeEthernet)
	ethernetPacket, _ := ethernetLayer.(*layers.Ethernet)

	data["SrcMAC"] = ethernetPacket.SrcMAC.String()
	data["DstMAC"] = ethernetPacket.DstMAC.String()
	data["NetworkProtocol"] = ethernetPacket.EthernetType.String()

	// 网络层
	if ethernetPacket.EthernetType.String() == "IPv4" {
		ipLayer := packet.Layer(layers.LayerTypeIPv4)
		ip, _ := ipLayer.(*layers.IPv4)

		data["IHL"] = int(ip.IHL)
		data["TOS"] = int(ip.TOS)
		data["Id"] = int(ip.Id)
		data["Flags"] = ip.Flags.String()
		data["FragOffset"] = int(ip.FragOffset)
		data["TTL"] = int(ip.TTL)
		data["SrcIP"] = ip.SrcIP.String()
		data["DstIP"] = ip.DstIP.String()
		data["TransportProtol"] = ip.NextLayerType().String()

	}

	// 传输层
	tcpLayer := packet.Layer(layers.LayerTypeTCP)
	if tcpLayer != nil {
		tcp, _ := tcpLayer.(*layers.TCP)
		data["SrcPort"] = int(tcp.SrcPort)
		data["DstPort"] = int(tcp.DstPort)

		if name, ok := layers.TCPPortNames[tcp.SrcPort]; ok {
			data["ApplicationProtocol"] = name
		}

		if name, ok := layers.TCPPortNames[tcp.DstPort]; ok {
			data["ApplicationProtocol"] = name
		}

		data["Seq"] = int64(tcp.Seq)
		data["Ack"] = int64(tcp.Ack)
		data["FIN"] = tcp.FIN
		data["SYN"] = tcp.SYN
		data["RST"] = tcp.RST
		data["PSH"] = tcp.PSH
		data["ACK"] = tcp.ACK
		data["URG"] = tcp.URG
		data["ECE"] = tcp.ECE
		data["CWR"] = tcp.CWR
		data["NS"] = tcp.NS
		data["Window"] = int(tcp.Window)
		data["CheckSum"] = int(tcp.Checksum)

		// for any, _ := range data {
		// 	fmt.Println(any, reflect.TypeOf(data[any]), data[any])
		// }
	}
	return data
}
